package crypter

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"os"
	"strings"

	"gitee.com/gousing/helper/internal"
)

/*
------------------------
RSA 加密解密
RSA 的加密机制有两种方案: RSA-OAEP 和 RSA-PKCS1-v1_5
推荐在新的应用中使用 RSA-OAEP
------------------------
RSA 签名
RSA 的签名机制机有种方案：RSA-PSS 和 RSA-PKCS1-v1_5
推荐在新的应用中使用 RSA-PSS
*/
const (
	RsaMakeKeysBitDefault int = 1024 // 密钥的位数: 512/1024/2048/4096...默认1024
)

var (
	ErrRsaDataEmpty       = errors.New("RSA data cannot be empty")
	ErrRsaBase64Decode    = errors.New("RSA data base64 decode error")
	ErrRsaPublicKeyEmpty  = errors.New("RSA publicKey cannot be empty")
	ErrRsaPrivateKeyEmpty = errors.New("RSA privateKey cannot be empty")
)

type RsaCrypto struct {
	publicBytes  []byte          //当前公钥字节符
	privateBytes []byte          //当前私钥字节符
	publicKey    *rsa.PublicKey  //当前公钥
	privateKey   *rsa.PrivateKey //当前私钥
}

func Rsa() *RsaCrypto {
	return &RsaCrypto{}
}

// SetPublicKey 设置公钥(支持PEM格式, 兼容PKCS#8和PKCS#1标准)
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func (r *RsaCrypto) SetPublicKey(publicKey []byte) (err error) {
	if len(publicKey) == 0 {
		return ErrRsaPublicKeyEmpty
	}
	r.publicBytes = publicKey
	r.publicKey, err = ParsePublicKey(r.publicBytes)
	return err
}

// SetPublicKeyWithString 设置公钥(支持PEM格式, 兼容PKCS#8和PKCS#1标准)
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func (r *RsaCrypto) SetPublicKeyWithString(publicKey string) (err error) {
	if len(publicKey) == 0 {
		return ErrRsaPublicKeyEmpty
	}
	r.publicBytes = internal.StringToBytes(publicKey)
	r.publicKey, err = ParsePublicKey(r.publicBytes)
	return err
}

// SetPublicKeyWithFile 设置公钥(支持PEM格式, 兼容PKCS#8和PKCS#1标准)
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func (r *RsaCrypto) SetPublicKeyWithFile(filePath string) (err error) {
	var publicKey []byte
	publicKey, err = os.ReadFile(filePath)
	if err != nil {
		return err
	}
	if len(publicKey) == 0 {
		return ErrRsaPublicKeyEmpty
	}
	r.publicBytes = publicKey
	r.publicKey, err = ParsePublicKey(r.publicBytes)
	return err
}

// GetPublicKey 获取公钥
func (r *RsaCrypto) GetPublicKey() []byte {
	return r.publicBytes
}

// GetPublicKeyToString 获取公钥(字符串)
func (r *RsaCrypto) GetPublicKeyToString() string {
	return internal.BytesToString(r.publicBytes)
}

// SetPrivateKey 设置私钥(支持PEM格式, 兼容PKCS#8和PKCS#1标准)
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func (r *RsaCrypto) SetPrivateKey(privateKey []byte) (err error) {
	if len(privateKey) == 0 {
		return ErrRsaPrivateKeyEmpty
	}
	r.privateBytes = privateKey
	r.privateKey, err = ParsePrivateKey(r.privateBytes)
	return err
}

// SetPrivateKeyWithString 设置私钥(支持PEM格式, 兼容PKCS#8和PKCS#1标准)
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func (r *RsaCrypto) SetPrivateKeyWithString(privateKey string) (err error) {
	if len(privateKey) == 0 {
		return ErrRsaPrivateKeyEmpty
	}
	r.privateBytes = internal.StringToBytes(privateKey)
	r.privateKey, err = ParsePrivateKey(r.privateBytes)
	return err
}

// SetPrivateKeyWithFile 设置私钥(支持PEM格式, 兼容PKCS#8和PKCS#1标准)
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func (r *RsaCrypto) SetPrivateKeyWithFile(file1Path string) (err error) {
	var privateKey []byte
	privateKey, err = os.ReadFile(file1Path)
	if err != nil {
		return err
	}
	if len(privateKey) == 0 {
		return ErrRsaPrivateKeyEmpty
	}
	r.privateBytes = privateKey
	r.privateKey, err = ParsePrivateKey(r.privateBytes)
	return err
}

// GetPrivateKey 获取私钥
func (r *RsaCrypto) GetPrivateKey() []byte {
	return r.privateBytes
}

// GetPrivateKeyToString 获取私钥(字符串)
func (r *RsaCrypto) GetPrivateKeyToString() string {
	return internal.BytesToString(r.privateBytes)
}

// Encrypt 公钥加密 (PKCS 标准)
//   - 遵循 PKCS#1 v1.5 标准(兼容性较好，在一些旧的系统或者对安全性要求不是极高的场景下使用)
func (r *RsaCrypto) Encrypt(data []byte) ([]byte, error) {
	if r.publicKey == nil {
		return []byte{}, ErrRsaPublicKeyEmpty
	}
	if len(data) == 0 {
		return []byte{}, ErrRsaDataEmpty
	}
	return rsa.EncryptPKCS1v15(rand.Reader, r.publicKey, data)
}

// EncryptBase64 公钥加密 (PKCS 标准), 加密结果使用 base64 编码后返回
//   - 遵循 PKCS#1 v1.5 标准(兼容性较好，在一些旧的系统或者对安全性要求不是极高的场景下使用)
func (r *RsaCrypto) EncryptBase64(data []byte) (string, error) {
	out, err := r.Encrypt(data)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(out), nil
}

// EncryptOAEP 公钥加密(OAEP 标准)
//   - 遵循 OAEP 标准(广泛应用于对安全性要求较高的场景，如金融交易、数字签名认证、安全通信协议, 如 SSL/TLS 在某些模式下)
//   - hashType 散列算法类型
//   - labels参数为可选参数, label用于与明文消息一起进行混淆操作, 加密时指定加密标签, 解密时用于验证标签是否一致
func (r *RsaCrypto) EncryptOAEP(data []byte, hashType crypto.Hash, labels ...[]byte) ([]byte, error) {
	if r.publicKey == nil {
		return []byte{}, ErrRsaPublicKeyEmpty
	}
	if len(data) == 0 {
		return []byte{}, ErrRsaDataEmpty
	}
	if !hashType.Available() {
		return []byte{}, errors.New("hashType not available")
	}
	var label []byte = nil
	if len(labels) > 0 {
		label = labels[0]
	}
	return rsa.EncryptOAEP(
		hashType.New(),
		rand.Reader,
		r.publicKey,
		data,
		label,
	)
}

// EncryptOAEPBase64 公钥加密(OAEP 标准), 加密结果使用 base64 编码后返回
//   - 遵循 OAEP 标准(广泛应用于对安全性要求较高的场景，如金融交易、数字签名认证、安全通信协议, 如 SSL/TLS 在某些模式下)
//   - hashType 散列算法类型
//   - labels参数为可选参数, label用于与明文消息一起进行混淆操作, 加密时指定加密标签, 解密时用于验证标签是否一致
func (r *RsaCrypto) EncryptOAEPBase64(data []byte, hashType crypto.Hash, labels ...[]byte) (string, error) {
	out, err := r.EncryptOAEP(data, hashType, labels...)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(out), nil
}

// Decrypt 私钥解密(PKCS标准)
//   - 遵循 PKCS#1 v1.5 标准(兼容性较好，在一些旧的系统或者对安全性要求不是极高的场景下使用)
func (r *RsaCrypto) Decrypt(data []byte) ([]byte, error) {
	if r.privateKey == nil {
		return []byte{}, ErrRsaPrivateKeyEmpty
	}
	if len(data) == 0 {
		return []byte{}, ErrRsaDataEmpty
	}
	return rsa.DecryptPKCS1v15(rand.Reader, r.privateKey, data)
}

// DecryptBase64 私钥解密(PKCS标准), 待解密的数据为 base64 编码的字符串
//   - 遵循 PKCS#1 v1.5 标准(兼容性较好，在一些旧的系统或者对安全性要求不是极高的场景下使用)
func (r *RsaCrypto) DecryptBase64(data string) ([]byte, error) {
	if r.privateKey == nil {
		return []byte{}, ErrRsaPrivateKeyEmpty
	}
	data = strings.TrimSpace(data)
	if len(data) == 0 {
		return []byte{}, ErrRsaDataEmpty
	}
	dataByte, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return []byte{}, ErrRsaBase64Decode
	}
	return r.Decrypt(dataByte)
}

// DecryptOAEP 私钥解密(OAEP 标准)
//   - 遵循 OAEP 标准(广泛应用于对安全性要求较高的场景，如金融交易、数字签名认证、安全通信协议, 如 SSL/TLS 在某些模式下)
//   - hashType 散列算法类型
//   - labels参数为可选参数, label用于与明文消息一起进行混淆操作, 加密时指定加密标签, 解密时用于验证标签是否一致
func (r *RsaCrypto) DecryptOAEP(data []byte, hashType crypto.Hash, labels ...[]byte) ([]byte, error) {
	if r.privateKey == nil {
		return []byte{}, ErrRsaPrivateKeyEmpty
	}
	if len(data) == 0 {
		return []byte{}, ErrRsaDataEmpty
	}
	if !hashType.Available() {
		return []byte{}, errors.New("hashType not available")
	}
	var label []byte = nil
	if len(labels) > 0 {
		label = labels[0]
	}
	return rsa.DecryptOAEP(hashType.New(), rand.Reader, r.privateKey, data, label)
}

// DecryptOAEPBase64 私钥解密(OAEP 标准), 待解密的数据为 base64 编码的字符串
//   - 遵循 OAEP 标准(广泛应用于对安全性要求较高的场景，如金融交易、数字签名认证、安全通信协议, 如 SSL/TLS 在某些模式下)
//   - hashType 散列算法类型
//   - labels参数为可选参数, label用于与明文消息一起进行混淆操作, 加密时指定加密标签, 解密时用于验证标签是否一致
func (r *RsaCrypto) DecryptOAEPBase64(data string, hashType crypto.Hash, labels ...[]byte) ([]byte, error) {
	if r.privateKey == nil {
		return []byte{}, ErrRsaPrivateKeyEmpty
	}
	data = strings.TrimSpace(data)
	if len(data) == 0 {
		return []byte{}, ErrRsaDataEmpty
	}
	dataByte, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return []byte{}, ErrRsaBase64Decode
	}
	return r.DecryptOAEP(dataByte, hashType, labels...)
}

// Sign (PKCS标准) 使用指定算法进行签名(私匙签名 Base64)
//   - 遵循 PKCS#1 v1.5 标准(兼容性较好，在一些旧的系统或者对安全性要求不是极高的场景下使用)
//   - hashType 散列算法类型
//   - 秘钥格式支持 PKCS#1 或 PKCS#8 格式
//   - 返回签名结果 Base64格式的字符串
func (r *RsaCrypto) Sign(data []byte, hashType crypto.Hash) (string, error) {
	if r.privateKey == nil {
		return "", ErrRsaPrivateKeyEmpty
	}
	if len(data) == 0 {
		return "", ErrRsaDataEmpty
	}
	if !hashType.Available() {
		return "", errors.New("hashType not available")
	}
	hashModel := hashType.New()
	hashModel.Write(data)
	signByte, err := rsa.SignPKCS1v15(rand.Reader, r.privateKey, hashType, hashModel.Sum(nil))
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(signByte), err
}

// VerifySign (PKCS标准) 使用指定算法验证签名(公匙验签 Base64)
//   - 遵循 PKCS#1 v1.5 标准(兼容性较好，在一些旧的系统或者对安全性要求不是极高的场景下使用)
//   - hashType 散列算法类型
//   - 待验证签名是 Base64格式的字符串
//   - 秘钥格式支持 PKCS#1 或 PKCS#8 格式
//   - 返回验证结果, 为nil表示验证成功
func (r *RsaCrypto) VerifySign(data []byte, hashType crypto.Hash, signature string) error {
	if r.privateKey == nil {
		return ErrRsaPrivateKeyEmpty
	}
	if len(data) == 0 {
		return ErrRsaDataEmpty
	}
	signature = strings.TrimSpace(signature)
	if len(signature) == 0 {
		return ErrRsaDataEmpty
	}
	signByte, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return ErrRsaBase64Decode
	}
	if !hashType.Available() {
		return errors.New("hashType not available")
	}
	hashModel := hashType.New()
	hashModel.Write(data)
	return rsa.VerifyPKCS1v15(r.publicKey, hashType, hashModel.Sum(nil), signByte)
}

// SignPSS (SignPSS标准) 使用指定算法进行签名(私匙签名 Base64)
//   - hashType 散列算法类型
//   - 秘钥格式支持 PKCS#1 或 PKCS#8 格式
//   - 返回签名结果 Base64格式的字符串
func (r *RsaCrypto) SignPSS(data []byte, hashType crypto.Hash) (string, error) {
	if r.privateKey == nil {
		return "", ErrRsaPrivateKeyEmpty
	}
	if len(data) == 0 {
		return "", ErrRsaDataEmpty
	}
	if !hashType.Available() {
		return "", errors.New("hashType not available")
	}
	hashModel := hashType.New()
	hashModel.Write(data)
	signByte, err := rsa.SignPSS(rand.Reader, r.privateKey, hashType, hashModel.Sum(nil), nil)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(signByte), err
}

// VerifySignPSS (SignPSS标准) 使用指定算法验证签名(公匙验签 Base64)
//   - hashType 散列算法类型
//   - 待验证签名是 Base64格式的字符串
//   - 秘钥格式支持 PKCS#1 或 PKCS#8 格式
//   - 返回验证结果, 为nil表示验证成功
func (r *RsaCrypto) VerifySignPSS(data []byte, hashType crypto.Hash, signature string) error {
	if r.privateKey == nil {
		return ErrRsaPrivateKeyEmpty
	}
	if len(data) == 0 {
		return ErrRsaDataEmpty
	}
	signature = strings.TrimSpace(signature)
	if len(signature) == 0 {
		return ErrRsaDataEmpty
	}
	signByte, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return ErrRsaBase64Decode
	}
	if !hashType.Available() {
		return errors.New("hashType not available")
	}
	hashModel := hashType.New()
	hashModel.Write(data)
	return rsa.VerifyPSS(r.publicKey, hashType, hashModel.Sum(nil), signByte, nil)
}

// MakeKeys 随机生成一对指定位数的RSA私钥和公钥(PEM格式/PKCS#8标准)
//   - 生成密钥的位数: bits => 512/1024/2048/4096...默认1024
//   - 密钥格式: PEM
//   - 密钥标准: PKCS#8
func (r *RsaCrypto) MakeKeys(bits ...int) (PrivateBytes []byte, PublicBytes []byte, err error) {
	bit := RsaMakeKeysBitDefault
	if len(bits) > 0 {
		bit = bits[0]
	}
	// 生成RSA密钥对
	// GenerateKey函数使用随机数据生成器random生成一对具有指定字位数的RSA密钥
	// Reader是一个全局、共享的密码用强随机数生成器

	var privateKey *rsa.PrivateKey
	privateKey, err = rsa.GenerateKey(rand.Reader, bit)
	if err != nil {
		return
	}
	// 将私钥转换为DER格式
	// 通过x509标准将得到的ras私钥序列化为ASN.1 的 DER编码
	var privateDer []byte
	privateDer, err = x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		return
	}
	// 处理私钥
	// PEM格式编码, 创建一个pem.Block结构体对象
	// PKCS#8 格式/私钥 BlockType为"PRIVATE KEY"
	privateBlock := pem.Block{Type: "PRIVATE KEY", Bytes: privateDer}
	// 将私钥转换为PEM格式
	PrivateBytes = pem.EncodeToMemory(&privateBlock)

	// 处理公钥
	publicKey := privateKey.PublicKey
	// 将公钥转换为DER格式
	// X509对公钥编码
	var PublicDer []byte
	PublicDer, err = x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		return
	}
	// PEM格式编码, 创建一个pem.Block结构体对象
	// PKCS#8 格式/公钥 BlockType为"PUBLIC KEY"
	publicBlock := pem.Block{Type: "PUBLIC KEY", Bytes: PublicDer}
	// 将公钥转换为PEM格式
	PublicBytes = pem.EncodeToMemory(&publicBlock)
	return
}

// MakeKeysPCKS1 随机生成一对指定位数的RSA私钥和公钥(PEM格式/PKCS#1标准)
//   - 生成密钥的位数: bits => 512/1024/2048/4096...默认1024
//   - 密钥格式: PEM
//   - 密钥标准: PKCS#1
func (r *RsaCrypto) MakeKeysPCKS1(bits ...int) (PrivateBytes []byte, PublicBytes []byte, err error) {
	bit := RsaMakeKeysBitDefault
	if len(bits) > 0 {
		bit = bits[0]
	}
	//GenerateKey函数使用随机数据生成器random生成一对具有指定字位数的RSA密钥
	//Reader是一个全局、共享的密码用强随机数生成器
	var privateKey *rsa.PrivateKey
	privateKey, err = rsa.GenerateKey(rand.Reader, bit)
	if err != nil {
		return
	}

	// 将私钥转换为DER格式
	// 通过x509标准将得到的ras私钥序列化为ASN.1 的 DER编码
	privateDer := x509.MarshalPKCS1PrivateKey(privateKey)

	// 处理私钥
	// PEM格式编码, 创建一个pem.Block结构体对象
	// PKCS#1 格式/私钥 BlockType为"RSA PRIVATE KEY"
	privateBlock := pem.Block{Type: "RSA PRIVATE KEY", Bytes: privateDer}
	// 将私钥转换为PEM格式
	PrivateBytes = pem.EncodeToMemory(&privateBlock)

	// 处理公钥
	publicKey := privateKey.PublicKey
	// 将公钥转换为DER格式
	// X509对公钥编码
	PublicDer := x509.MarshalPKCS1PublicKey(&publicKey)
	// PEM格式编码, 创建一个pem.Block结构体对象
	// PKCS#1 格式/公钥 BlockType为"RSA PUBLIC KEY"
	publicBlock := pem.Block{Type: "RSA PUBLIC KEY", Bytes: PublicDer}
	//保存
	PublicBytes = pem.EncodeToMemory(&publicBlock)
	return
}

// MakeKeysToString 随机生成一对指定位数的RSA私钥和公钥(PEM格式/PKCS#8标准)
//   - 生成密钥的位数: bits => 512/1024/2048/4096...默认1024
//   - 密钥格式: PEM
//   - 密钥标准: PKCS#8
func (r *RsaCrypto) MakeKeysToString(bits ...int) (PrivateKey string, PublicKey string, err error) {
	var (
		PrivateBytes []byte
		PublicBytes  []byte
	)
	PrivateBytes, PublicBytes, err = r.MakeKeys(bits...)
	if err != nil {
		return
	}
	PrivateKey = internal.BytesToString(PrivateBytes)
	PublicKey = internal.BytesToString(PublicBytes)
	return
}

// MakeKeysPKCS1ToString 随机生成一对指定位数的RSA私钥和公钥(PEM格式/PKCS#1标准)
//   - 生成密钥的位数: bits => 512/1024/2048/4096...默认1024
//   - 密钥格式: PEM
//   - 密钥标准: PKCS#1
func (r *RsaCrypto) MakeKeysPKCS1ToString(bits ...int) (PrivateKey string, PublicKey string, err error) {
	var (
		PrivateBytes []byte
		PublicBytes  []byte
	)
	PrivateBytes, PublicBytes, err = r.MakeKeysPCKS1(bits...)
	if err != nil {
		return
	}
	PrivateKey = internal.BytesToString(PrivateBytes)
	PublicKey = internal.BytesToString(PublicBytes)
	return
}

// ParsePublicKey 转化公钥为 rsa.PublicKey 对象
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func ParsePublicKey(publicKey []byte) (*rsa.PublicKey, error) {
	// decode public key
	block, _ := pem.Decode(publicKey)
	if block == nil {
		return nil, errors.New("ParsePublicKey pem.Decode解析失败")
	}
	// block.Type 值通用为大写表示 RSA PUBLIC KEY  / PUBLIC KEY
	blockType := strings.ToUpper(block.Type)
	if blockType == "PUBLIC KEY" {
		// PCKS#8 格式
		pub, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, errors.New("ParsePublicKey[ blockType:" + blockType + "] 解析失败, " + err.Error())
		}
		if obj, ok := pub.(*rsa.PublicKey); ok {
			return obj, nil
		}
		return nil, errors.New("ParsePublicKey[ blockType:" + blockType + "] 解析成功, 转换为 *rsa.PublicKey 失败")
	} else if blockType == "RSA PUBLIC KEY" {
		// PCKS#1 格式
		pub, err := x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return nil, errors.New("ParsePublicKey[ blockType:" + blockType + "] 解析失败, " + err.Error())
		}
		return pub, nil
	}
	return nil, errors.New("ParsePublicKey[ blockType:" + blockType + "] 未知的公钥类型, 对于String或[]Byte格式的证书仅支持PEM格式")
}

// ParsePrivateKey 转化私钥为rsa.PrivateKey对象
//   - 支持 PEM 格式
//   - 兼容 PKCS#1 或 PKCS#8 标准
func ParsePrivateKey(privateKey []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(privateKey)
	if block == nil {
		return nil, errors.New("ParsePrivateKey pem.Decode解析失败")
	}
	// block.Type 值通用为大写表示 RSA PRIVATE KEY  / PRIVATE KEY
	// 也有大小写混合方式 RSA Private Key /  Private Key
	blockType := strings.ToUpper(block.Type)
	if blockType == "RSA PRIVATE KEY" {
		// 尝试 PKCS#1 格式解析
		keyPKCS1, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		return keyPKCS1, nil
	} else if blockType == "PRIVATE KEY" {
		// 尝试 PKCS#8 格式解析
		keyPKCS8, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		if obj, ok := keyPKCS8.(*rsa.PrivateKey); ok {
			return obj, nil
		}
		return nil, errors.New("ParsePrivateKey x509解析成功, 转换为*rsa.PrivateKey失败")
	}
	return nil, errors.New("ParsePrivateKey 未知的私钥类型[" + block.Type + "], 对于String或[]Byte格式的证书仅支持PEM格式")
}
